package lzw;

import dcc.CompressionAlgorithm;
import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;


public class LZW  implements CompressionAlgorithm{
        public static final int INITIAL_DICT_SIZE = 256;
        public static final int DEFAULT_CODEWORD_LENGTH = 12;

        private int codeWordLength = DEFAULT_CODEWORD_LENGTH;
        private Map<ByteArray, Integer> codeTable;
        private List<ByteArray> decodeTable;
        
        @Override
        public void compress(InputStream in, OutputStream out) throws IOException {
            //initialize dictionary [1]
            init();
            int code = INITIAL_DICT_SIZE;
            
            InputStream bufferedIn = new BufferedInputStream(in);
            Output compressedOutput = new Output(new BufferedOutputStream(out),codeWordLength);
            //read first byte [2]
            int firstByte = bufferedIn.read();
            ByteArray w = new ByteArray((byte) firstByte);
            int k;
            //continue reading to the end of the file [2]
            while ((k = bufferedIn.read()) != -1) {
                //combine the new symbol with previously read [3]
                ByteArray wK = new ByteArray(w).append((byte) k);
                //if it is already in the table, change the previously read symbol [3]
                if (codeTable.containsKey(wK)) {
                    w = wK;
                } else {
                    //if it's not, write the appropriate symbol to the output file (index of the matching sequence) [4]
                    compressedOutput.write(codeTable.get(w));
                    if (code < (1 << codeWordLength) - 1) {
                        //add symbol to the dictionary [5]
                        codeTable.put(wK, code++);
                    }
                    //change the previously read symbol [6]
                    w = new ByteArray((byte) k);
                }
            }
            //write the last symbol to the stream
            compressedOutput.write(codeTable.get(w));
            compressedOutput.flush();
            compressedOutput.close();
        }

        @Override
        public void decompress(InputStream in, OutputStream out) throws IOException {
            //initialize the alphabet [1]
            init();            
            Input compressedInput = new Input(new BufferedInputStream(in),codeWordLength);
            OutputStream bufferedOut = new BufferedOutputStream(out);
            //read 1 code from the input [2]
            int oldCode = compressedInput.read();
            //write that symbol [3]
            bufferedOut.write(oldCode);
            int character = oldCode;
            int newCode;
            //read the compressed file to the end [4]
            while((newCode = compressedInput.read()) != -1) {
                ByteArray string;
                
                if(newCode >= decodeTable.size()) {
                    string = new ByteArray(decodeTable.get(oldCode));
                    string.append((byte) character);
                } else {
                    string = decodeTable.get(newCode);
                }
                //output the resulting ByteArray [6]
                for(int i = 0; i < string.size(); i++) {
                    bufferedOut.write(string.get(i));
                }
                character = string.get(0);
                decodeTable.add(new ByteArray(decodeTable.get(oldCode)).append((byte) character));
                oldCode = newCode;
            }
            bufferedOut.flush();
            bufferedOut.close();
        }
        
        private void init() {
            codeTable = new HashMap<>();
            decodeTable = new ArrayList<>();
            for (int i = 0; i < INITIAL_DICT_SIZE; i++) {
                codeTable.put(new ByteArray((byte) i), i);
                decodeTable.add(new ByteArray((byte) i));
            }
        }
}
